"""
数据备份与恢复工具
用于管理数据的定期备份和灾难恢复
"""
import os
import time
import shutil
import tarfile
import datetime
from typing import List, Optional
from fin_senti_entity_platform.utils.logger import Logger
from fin_senti_entity_platform.utils.config_loader import ConfigLoader
from fin_senti_entity_platform.data_collection.storage.storage_proxy import storage_proxy


class DataBackup:
    """数据备份与恢复类"""
    
    def __init__(self):
        """初始化数据备份与恢复工具"""
        self.config = ConfigLoader().get('data_backup', {})
        self.logger = Logger().get_logger(__name__)
        self.backup_dir = self.config.get('backup_dir', './data/backups')
        self.retention_days = self.config.get('retention_days', 7)
        
        # 确保备份目录存在
        os.makedirs(self.backup_dir, exist_ok=True)
    
    def _get_backup_filename(self, prefix: str = 'backup') -> str:
        """
        生成备份文件名
        
        Args:
            prefix: 文件名前缀
        
        Returns:
            str: 备份文件名
        """
        timestamp = datetime.datetime.now().strftime('%Y%m%d_%H%M%S')
        return f"{prefix}_{timestamp}.tar.gz"
    
    def backup_mongodb(self, db_name: str, collection_names: Optional[List[str]] = None) -> str:
        """
        备份MongoDB数据库
        
        Args:
            db_name: 数据库名称
            collection_names: 要备份的集合列表，如果为None则备份所有集合
        
        Returns:
            str: 备份文件路径
        """
        try:
            self.logger.info(f"开始备份MongoDB数据库: {db_name}")
            
            # 获取MongoDB存储实例
            mongo_storage = storage_proxy.get_storage('mongodb')
            if not mongo_storage:
                self.logger.error("MongoDB存储未初始化")
                return ""
                
            # 确保临时备份目录存在
            temp_backup_dir = os.path.join(self.backup_dir, f"temp_{int(time.time())}")
            os.makedirs(temp_backup_dir, exist_ok=True)
            
            # 获取所有集合名称
            if collection_names is None:
                collection_names = mongo_storage.db.list_collection_names()
                
            # 备份每个集合
            for collection_name in collection_names:
                self.logger.info(f"正在备份集合: {collection_name}")
                
                # 查询集合中的所有数据
                all_data = mongo_storage.query(collection_name, {})
                
                if all_data:
                    # 保存为JSON文件
                    import json
                    json_file = os.path.join(temp_backup_dir, f"{collection_name}.json")
                    with open(json_file, 'w', encoding='utf-8') as f:
                        json.dump(all_data, f, ensure_ascii=False, indent=2)
                        
            # 创建压缩文件
            backup_filename = self._get_backup_filename(f"mongodb_{db_name}")
            backup_path = os.path.join(self.backup_dir, backup_filename)
            
            with tarfile.open(backup_path, "w:gz") as tar:
                tar.add(temp_backup_dir, arcname=os.path.basename(temp_backup_dir))
                
            # 清理临时文件
            shutil.rmtree(temp_backup_dir)
            
            # 清理过期备份
            self._cleanup_old_backups()
            
            self.logger.info(f"MongoDB数据库备份成功: {backup_path}")
            return backup_path
            
        except Exception as e:
            self.logger.error(f"MongoDB数据库备份失败: {str(e)}")
            # 清理临时文件
            if 'temp_backup_dir' in locals() and os.path.exists(temp_backup_dir):
                shutil.rmtree(temp_backup_dir)
            return ""
    
    def restore_mongodb(self, backup_file: str, db_name: str, collection_names: Optional[List[str]] = None) -> bool:
        """
        从备份文件恢复MongoDB数据库
        
        Args:
            backup_file: 备份文件路径
            db_name: 数据库名称
            collection_names: 要恢复的集合列表，如果为None则恢复所有集合
        
        Returns:
            bool: 恢复是否成功
        """
        try:
            self.logger.info(f"开始从备份文件恢复MongoDB数据库: {backup_file}")
            
            # 检查备份文件是否存在
            if not os.path.exists(backup_file):
                self.logger.error(f"备份文件不存在: {backup_file}")
                return False
                
            # 获取MongoDB存储实例
            mongo_storage = storage_proxy.get_storage('mongodb')
            if not mongo_storage:
                self.logger.error("MongoDB存储未初始化")
                return False
                
            # 创建临时恢复目录
            temp_restore_dir = os.path.join(self.backup_dir, f"temp_restore_{int(time.time())}")
            os.makedirs(temp_restore_dir, exist_ok=True)
            
            # 解压备份文件
            with tarfile.open(backup_file, "r:gz") as tar:
                tar.extractall(path=temp_restore_dir)
                
            # 获取解压后的目录
            extracted_dir = os.path.join(temp_restore_dir, os.listdir(temp_restore_dir)[0])
            
            # 恢复每个集合
            for filename in os.listdir(extracted_dir):
                if filename.endswith('.json'):
                    collection_name = filename[:-5]  # 移除.json后缀
                    
                    # 检查是否在要恢复的集合列表中
                    if collection_names is None or collection_name in collection_names:
                        self.logger.info(f"正在恢复集合: {collection_name}")
                        
                        # 读取JSON文件
                        json_file = os.path.join(extracted_dir, filename)
                        import json
                        with open(json_file, 'r', encoding='utf-8') as f:
                            data = json.load(f)
                            
                        # 清空集合并恢复数据
                        mongo_storage.delete(collection_name, {}, multi=True)
                        if data:
                            mongo_storage.batch_save(data, collection_name)
                            
            # 清理临时文件
            shutil.rmtree(temp_restore_dir)
            
            self.logger.info(f"MongoDB数据库恢复成功")
            return True
            
        except Exception as e:
            self.logger.error(f"MongoDB数据库恢复失败: {str(e)}")
            # 清理临时文件
            if 'temp_restore_dir' in locals() and os.path.exists(temp_restore_dir):
                shutil.rmtree(temp_restore_dir)
            return False
    
    def backup_parquet_files(self, source_dir: str = './data/parquet') -> str:
        """
        备份Parquet文件
        
        Args:
            source_dir: Parquet文件所在目录
        
        Returns:
            str: 备份文件路径
        """
        try:
            self.logger.info(f"开始备份Parquet文件: {source_dir}")
            
            # 检查源目录是否存在
            if not os.path.exists(source_dir):
                self.logger.error(f"Parquet文件目录不存在: {source_dir}")
                return ""
                
            # 创建压缩文件
            backup_filename = self._get_backup_filename("parquet")
            backup_path = os.path.join(self.backup_dir, backup_filename)
            
            with tarfile.open(backup_path, "w:gz") as tar:
                tar.add(source_dir, arcname=os.path.basename(source_dir))
                
            # 清理过期备份
            self._cleanup_old_backups()
            
            self.logger.info(f"Parquet文件备份成功: {backup_path}")
            return backup_path
            
        except Exception as e:
            self.logger.error(f"Parquet文件备份失败: {str(e)}")
            return ""
    
    def restore_parquet_files(self, backup_file: str, target_dir: str = './data/parquet') -> bool:
        """
        从备份文件恢复Parquet文件
        
        Args:
            backup_file: 备份文件路径
            target_dir: 恢复目标目录
        
        Returns:
            bool: 恢复是否成功
        """
        try:
            self.logger.info(f"开始从备份文件恢复Parquet文件: {backup_file}")
            
            # 检查备份文件是否存在
            if not os.path.exists(backup_file):
                self.logger.error(f"备份文件不存在: {backup_file}")
                return False
                
            # 确保目标目录存在
            os.makedirs(target_dir, exist_ok=True)
            
            # 解压备份文件
            with tarfile.open(backup_file, "r:gz") as tar:
                tar.extractall(path=target_dir)
                
            self.logger.info(f"Parquet文件恢复成功")
            return True
            
        except Exception as e:
            self.logger.error(f"Parquet文件恢复失败: {str(e)}")
            return False
    
    def _cleanup_old_backups(self):
        """清理过期的备份文件"""
        try:
            self.logger.info(f"开始清理 {self.retention_days} 天前的过期备份")
            
            # 计算过期时间
            cutoff_time = time.time() - (self.retention_days * 24 * 60 * 60)
            
            # 遍历备份目录中的所有文件
            for filename in os.listdir(self.backup_dir):
                file_path = os.path.join(self.backup_dir, filename)
                
                # 检查文件是否为备份文件（以.tar.gz结尾）
                if os.path.isfile(file_path) and file_path.endswith('.tar.gz'):
                    # 获取文件的修改时间
                    file_mtime = os.path.getmtime(file_path)
                    
                    # 如果文件过期，则删除
                    if file_mtime < cutoff_time:
                        os.remove(file_path)
                        self.logger.info(f"已删除过期备份文件: {filename}")
                        
        except Exception as e:
            self.logger.error(f"清理过期备份文件失败: {str(e)}")


# 创建全局数据备份实例
backup_manager = DataBackup()